PyBaobabDT

Description: The pybaobabdt package provides a python implementation for the visualization of decision trees. The technique is based on the scientific paper BaobabView: Interactive construction and analysis of decision trees developed by the TU/e. https://pypi.org/project/pybaobabdt/

import pybaobabdt
import pandas as pd
from scipy.io import arff
from sklearn.tree import DecisionTreeClassifier

Description: Due to incompatibility with our catboost model, a scikit model trained on sample data (taken from https://pypi.org/project/pybaobabdt/ ) to visualize decision tree was used.

data = arff.loadarff('data/winequality-red.arff')
df   = pd.DataFrame(data[0])

y = list(df['class'])
features = list(df.columns)
features.remove('class')
X = df.loc[:, features]

clf = DecisionTreeClassifier().fit(X,y);

ax = pybaobabdt.drawTree(clf, size=10, dpi=72, features=features)

<Figure size 720x720 with 0 Axes>

Description: Visualization shows data streams flowing across each node displaying concerned split feature and threshold (each class is represented by a color, the width of the link represents the number of items flowing from one node to the other). Although the visualization solution provided by the pybaobab library fully meets the need to see data flows flowing through the model structure, it is not compatible with catboost models.

ax.get_figure().savefig('visualizations/tree.png', format='png', dpi=300, transparent=True)

Description: Also, trees from a RandomForest classifier can be visualized:

import matplotlib.pyplot as plt
from sklearn.ensemble import RandomForestClassifier

data = arff.loadarff('data/vehicle.arff')

df = pd.DataFrame(data[0])
y = list(df['class'])
features = list(df.columns)
features.remove('class')
X = df.loc[:, features]

clf = RandomForestClassifier(n_estimators=20, n_jobs=-1, random_state=0);
clf.fit(X, y);

size = (15,15)
plt.rcParams['figure.figsize'] = size
fig = plt.figure(figsize=size, dpi=300)

for idx, tree in enumerate(clf.estimators_):
    ax1 = fig.add_subplot(5, 4, idx+1)
    pybaobabdt.drawTree(tree, model=clf, size=15, dpi=300, features=features, ax=ax1)
    
fig.savefig('visualizations/random-forest.png', format='png', dpi=1200, transparent=True)